#!/usr/bin/env python
# coding:utf-8
from PoboAPI import *
import datetime
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

# 多因子选股策略示例
# 策略逻辑：结合动量、波动率和价值三个因子进行股票筛选和交易

def OnStart(context):
    print("多因子选股策略启动...")
    
    # 登录交易账号
    context.myacc = None
    if "回测股票" in context.accounts:
        print("登录交易账号[回测股票]")
        if context.accounts["回测股票"].Login():
            context.myacc = context.accounts["回测股票"]
    
    # 初始化参数
    g.universe = ['600000.SHSE', '600036.SHSE', '601318.SHSE', '600519.SHSE', '000001.SZSE', 
                 '000333.SZSE', '000651.SZSE', '000858.SZSE', '002415.SZSE', '300750.SZSE']  # 股票池
    g.stock_num = 3  # 持仓股票数量
    g.position_weight = 0.3  # 单只股票仓位比例
    g.period = 20  # 因子计算周期
    g.rebalance_days = 5  # 调仓周期(交易日)
    g.day_count = 0  # 当前交易日计数
    g.positions = []  # 当前持仓
    
    # 订阅行情
    for stock in g.universe:
        SubscribeBar(stock, BarType.Day)
    
    # 设置每日收盘前检查的闹钟
    g.alarm_id = SetAlarm(datetime.time(14, 50), RepeatType.Daily)

# 闹钟事件
def OnAlarm(context, alarmid):
    # 收盘前检查持仓
    print("收盘前检查持仓...")
    check_positions(context)

# 行情初始化事件
def OnMarketQuotationInitialEx(context, exchange, daynight):
    if daynight != 'day':
        return
    
    g.day_count += 1
    print(f"交易日: {g.day_count}")
    
    # 每隔rebalance_days天进行调仓
    if g.day_count % g.rebalance_days == 0:
        rebalance_portfolio(context)

# 计算动量因子
def calculate_momentum_factor(stock_data):
    # 计算过去20天的收益率
    returns = stock_data['close'].pct_change(g.period).iloc[-1]
    return returns

# 计算波动率因子
def calculate_volatility_factor(stock_data):
    # 计算过去20天的波动率，波动率越低分数越高
    volatility = stock_data['close'].pct_change().std() * np.sqrt(252)
    return -volatility  # 取负值，使得波动率低的股票得分高

# 计算价值因子
def calculate_value_factor(stock_data):
    # 这里使用收盘价与20日均线的比值作为简化的价值因子
    # 实际应用中可以使用PE、PB等基本面指标
    ma20 = stock_data['close'].rolling(window=20).mean().iloc[-1]
    value = ma20 / stock_data['close'].iloc[-1]
    return value

# 计算综合因子得分
def calculate_factor_score(context):
    factor_data = pd.DataFrame(index=g.universe)
    
    # 计算各个因子
    momentum_scores = []
    volatility_scores = []
    value_scores = []
    
    for stock in g.universe:
        try:
            # 获取历史数据
            df = GetHisDataAsDF(stock, bar_type=BarType.Day)
            if len(df) < g.period + 10:
                print(f"{stock} 历史数据不足，跳过")
                momentum_scores.append(np.nan)
                volatility_scores.append(np.nan)
                value_scores.append(np.nan)
                continue
                
            # 计算各因子得分
            momentum = calculate_momentum_factor(df)
            volatility = calculate_volatility_factor(df)
            value = calculate_value_factor(df)
            
            momentum_scores.append(momentum)
            volatility_scores.append(volatility)
            value_scores.append(value)
            
        except Exception as e:
            print(f"计算 {stock} 因子时出错: {str(e)}")
            momentum_scores.append(np.nan)
            volatility_scores.append(np.nan)
            value_scores.append(np.nan)
    
    # 将因子得分添加到DataFrame
    factor_data['momentum'] = momentum_scores
    factor_data['volatility'] = volatility_scores
    factor_data['value'] = value_scores
    
    # 标准化因子得分
    scaler = StandardScaler()
    numeric_data = factor_data.dropna()
    if len(numeric_data) > 0:
        scaled_data = scaler.fit_transform(numeric_data)
        scaled_df = pd.DataFrame(scaled_data, index=numeric_data.index, columns=numeric_data.columns)
        
        # 计算综合得分 (各因子权重可调整)
        scaled_df['total_score'] = scaled_df['momentum'] * 0.4 + scaled_df['volatility'] * 0.3 + scaled_df['value'] * 0.3
        
        # 按得分排序
        scaled_df = scaled_df.sort_values('total_score', ascending=False)
        return scaled_df
    else:
        print("没有足够的有效数据进行因子计算")
        return pd.DataFrame()

# 调仓函数
def rebalance_portfolio(context):
    print("开始调仓...")
    
    # 计算因子得分
    factor_scores = calculate_factor_score(context)
    if factor_scores.empty:
        print("无法获取有效的因子得分，跳过调仓")
        return
    
    # 获取得分最高的前N只股票
    top_stocks = factor_scores.index[:g.stock_num].tolist()
    print(f"选出的股票: {top_stocks}")
    
    # 获取当前持仓
    current_positions = get_positions(context)
    
    # 卖出不在目标持仓中的股票
    for stock in current_positions:
        if stock not in top_stocks:
            sell_stock(context, stock)
    
    # 买入新的目标股票
    for stock in top_stocks:
        if stock not in current_positions:
            buy_stock(context, stock)
    
    # 更新持仓记录
    g.positions = get_positions(context)
    print(f"调仓后持仓: {g.positions}")

# 获取当前持仓股票列表
def get_positions(context):
    positions = []
    if context.myacc is None:
        return positions
    
    option = PBObj()
    pos_list = context.myacc.GetPositions(option)
    for pos in pos_list:
        positions.append(pos.contract)
    
    return positions

# 买入股票
def buy_stock(context, stock):
    if context.myacc is None:
        print("交易账户未登录，无法交易")
        return
    
    try:
        # 获取最新价格
        quote = GetQuote(stock)
        price = quote.now
        
        # 计算买入数量
        cash = context.myacc.Cash
        target_value = cash * g.position_weight
        volume = int(target_value / price / 100) * 100  # 确保是100的整数倍
        
        if volume >= 100:  # 最小交易单位为100股
            price_type = PriceType(PbPriceType.Limit, 16, 0)
            QuickInsertOrder(context.myacc, stock, 'buy', 'open', price_type, volume)
            print(f"买入 {stock} {volume}股，价格 {price}")
        else:
            print(f"买入 {stock} 数量不足100股，跳过")
    except Exception as e:
        print(f"买入 {stock} 时出错: {str(e)}")

# 卖出股票
def sell_stock(context, stock):
    if context.myacc is None:
        print("交易账户未登录，无法交易")
        return
    
    try:
        # 获取持仓数量
        option = PBObj()
        option.contract = stock
        pos_list = context.myacc.GetPositions(option)
        
        for pos in pos_list:
            if pos.bstype.BuySellFlag == "0":  # 多头持仓
                price_type = PriceType(PbPriceType.Limit, 16, 0)
                QuickInsertOrder(context.myacc, stock, 'sell', 'close', price_type, pos.volume)
                print(f"卖出 {stock} {pos.volume}股")
    except Exception as e:
        print(f"卖出 {stock} 时出错: {str(e)}")

# 检查持仓
def check_positions(context):
    if context.myacc is None:
        return
    
    positions = get_positions(context)
    print(f"当前持仓: {positions}")
    
    # 计算持仓盈亏
    total_profit = 0
    for stock in positions:
        option = PBObj()
        option.contract = stock
        pos_list = context.myacc.GetPositions(option)
        
        for pos in pos_list:
            profit = pos.profit
            total_profit += profit
            print(f"{stock} 盈亏: {profit}")
    
    print(f"总盈亏: {total_profit}")

# K线事件
def OnBar(context, code, bartype):
    # 这里可以添加实时监控逻辑
    pass